[WIP] Use cupy, in which case all operations are performed on GPU
#259
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
lasy/__init__.py
Outdated
| @@ -1 +1,8 @@ | |||
| __version__ = "0.4.0" | |||
|
|
|||
There was a problem hiding this comment.
This should be discussed: right now, LASY automatically detects whether cupy is available and always uses it if it is.
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
RemiLehe
changed the title
cupy, in which case all operations are performed on GPU
for more information, see https://pre-commit.ci
| if self.dim == "rt": | ||
| # Construct the propagator (check if exists) | ||
| if not hasattr(self, "prop"): | ||
| spatial_axes = (self.grid.axes[0],) | ||
| self.prop = [] | ||
| if use_cupy: | ||
| # Move quantities to CPU to create propagator |
There was a problem hiding this comment.
@hightower8083 Would it be possible to modify the PropagatorResampling, so that it can take cupy arrays as input?
Right now, if don't move k and spatial_axes to the CPU, I get the following error:
lasy/laser.py:254: in propagate
PropagatorResampling(
.../python-3.11/lib/python3.11/site-packages/axiprop/lib.py:242: in __init__
self.init_kr(self.Rmax, self.Nr)
.../python-3.11/lib/python3.11/site-packages/axiprop/common.py:113: in init_kr
self.kr = self.alpha / Rmax
Note that I am using this branch of axiprop
There was a problem hiding this comment.
well there is no need to have k and spatial_axes on GPU -- spatial axes are used to make transverse wave-numbers k_r that are transferred to GPU by axiprop and k is used per-element in the loop as scalars
| @@ -737,9 +738,12 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"): | |||
| time_axis_indx = -1 | |||
|
|
|||
| t_axis = grid.axes[time_axis_indx] | |||
| if use_cupy: | |||
| t_axis = xp.asnumpy(t_axis) | |||
There was a problem hiding this comment.
If I don't do this, I later get:
lasy/utils/laser_utils.py:744: in export_to_z
FieldAxprp = ScalarFieldEnvelope(omega0 / c, t_axis)
.../python-3.11/lib/python3.11/site-packages/axiprop/containers.py:101: in __init__
self.k_freq_base = 2 * np.pi * np.fft.fftfreq(self.Nt, c*self.dt)
.../lib/python3.11/site-packages/numpy/fft/helper.py:169: in fftfreq
return results * val
cupy/_core/core.pyx:1697: in cupy._core.core._ndarray_base.__array_ufunc__
???
cupy/_core/_kernel.pyx:1283: in cupy._core._kernel.ufunc.__call__
???
cupy/_core/_kernel.pyx:159: in cupy._core._kernel._preprocess_args
???
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
> ???
E TypeError: Unsupported type <class 'numpy.ndarray'>
There was a problem hiding this comment.
t_axis is actually not needed on GPU as all propagation in done in the spectral domain -- as soon as we do FFT we don't have time. self.k_freq_base can remain on the host
for more information, see https://pre-commit.ci
| @@ -783,10 +787,10 @@ def export_to_z(dim, grid, omega0, z_axis=None, z0=0.0, t0=0.0, backend="NP"): | |||
| verbose=False, | |||
| ) | |||
| # Convert the spectral image to the spatial field representation | |||
| FieldAxprp.import_field(np.moveaxis(field, -1, 0).copy()) | |||
| FieldAxprp.import_field(xp.moveaxis(field, -1, 0).copy()) | |||
There was a problem hiding this comment.
I am currently getting an error of the form:
lasy/utils/laser_utils.py:790: in export_to_z
FieldAxprp.import_field(xp.moveaxis(field, -1, 0).copy())
.../python-3.11/lib/python3.11/site-packages/axiprop/containers.py:369: in import_field
self.time_to_frequency()
.../python-3.11/lib/python3.11/site-packages/axiprop/containers.py:439: in time_to_frequency
self.Field_ft[:] = np.fft.ifft(self.Field, axis=0, norm="backward")
There was a problem hiding this comment.
we should get rid of the axiprop Containers -- they currently only have CPU implementation for spectral transforms. I assume we are doing FFT now with lasy portable methods, so lets stick to this
| @@ -0,0 +1,10 @@ | |||
| try: | |||
There was a problem hiding this comment.
Re backend control: #259 (comment)
We could do a similar control as done in matplotlib, e.g.,
import lasy
lasy.use("numpy") # default: "auto"https://matplotlib.org/stable/users/explain/figure/backends.html
The logic could be, using the usual precedence of options:
- check if this option was set (on the module aka the current process), otherwise
- check env variable to use default, otherwise
- use cupy if found, otherwise
- use numpy
This PR adds a new file
backend.py, that defines a quantityxpto be:cupyifcupyinstallednumpyotherwiseAlmost all array operations in
lasythen usexp(e.g. to allocate the fields array, to perform the FFT, etc.)Note that, in the current state of the PR, the user has no control of whether
cupywill be used or not: as long ascupyis installed,lasywill automatically use it. At this point, there is also no message telling the user whethercupyornumpyis used (but the user can always printlasy.backend.use_cupy).The getters for
temporal_fieldandspectral_fieldhave also been adapted, so as to be able to get the fields on CPU. This is used e.g. for the functionsshow(plotting with matplotlib) andwrite_to_file(dumping to disk with openPMD), as well as in several tests.